import numpy as np
import scipy.constants as C
import scipy.special
import matplotlib.pyplot as plt
from common import  Ek_to_gamma

def calculate_E_sc_disk(z_i_list, Q, p_zi_list,a, b, N_zeros, z_range):
    """
    多个电子圆盘产生的空间电荷场
    :param z_i_list: 电子圆盘位置
    :param Q: 电子圆盘电荷量
    :param p_zi_list: 每个圆盘对应的电子动能(eV)
    :param a: 腔体半径
    :param b: 电子圆盘半径
    :param N_zeros: 使用的Bessel函数零点数目
    :param z_range: z轴上的采样点数组
    :return: 空间电荷场E_sc
    """
    # Bessel函数J0的前N_zeros个零点
    mu_0_ps = scipy.special.jn_zeros(0, N_zeros)

    # 计算单个电子圆盘 z_i 对应的电场分量
    def E_sc_disk(z, z_i, Q, p_zi, a, b):
        """
        单个电子圆盘z轴上的空间电荷场
        :param z:
        :param z_i: 电子圆盘位置
        :param Q: 电子圆盘电荷量
        :param p_zi: 电子圆盘动能
        :param a: 腔体半径
        :param b: 电子圆盘半径
        :return: 空间电荷场E_sc_disk
        """
        gamma = Ek_to_gamma(p_zi)

        def Ez_component(mu_n):  # 计算单个Bessel函数零点的电场分量
            arr = (
                np.exp(-mu_n * gamma * np.abs(z - z_i)) *
                (1 / mu_n * scipy.special.j1(mu_n * b / a) / scipy.special.j1(mu_n)) ** 2
            )
            return arr

        # 累加所有分量
        # arr_sum = np.sum((Ez_component(mu_n) for mu_n in mu_0_ps), axis=0)
        arr_sum = sum((Ez_component(mu_n) for mu_n in mu_0_ps))
        return Q / (2 * np.pi * C.epsilon_0 * b ** 2) * arr_sum * np.sign(z - z_i)

    # 计算所有z_i与对应γ对z_range的电场叠加
    E_sc = np.zeros_like(z_range)
    for z_i, p_zi in zip(z_i_list, p_zi_list):  # 对Z_i及该位置对应的动能Ek_eV同时遍历,不同速度粒子对应修正因子γ不同
        E_sc += E_sc_disk(z_range, z_i, Q, p_zi, a, b)

    return E_sc

# 单独测试
if __name__ == '__main__':
    plt.ion()
    z_i_list = [0, 2e-3]  # 单个圆盘位置
    Q = 1e-9  # 圆盘电荷量 (C)
    p_zi_list = [50e7, 20e7]  # 每个圆盘对应的电子动能
    a = 4e-3  # 漂移管半径 (m)
    b = 3e-3  # 电子圆盘半径 (m)
    N_zeros = 70  # 使用的 Bessel 函数零点数目
    z_range = np.linspace(-10e-3, 10e-3, 1000)  # z 轴上的采样点数组 (m)

    # 计算电场
    E_sc = calculate_E_sc_disk(z_i_list, Q, p_zi_list, a, b, N_zeros, z_range)

    # 绘制结果
    plt.figure()
    plt.plot(z_range * 1e3, E_sc, label=f"Ez, N_zeros={N_zeros}")
    plt.title(f"Distribution of Disk mode (Ez)")
    plt.xlabel("Distance of z (mm)")
    plt.ylabel("Electric Field Ez (V/m)")
    plt.legend()
    plt.grid()
    plt.show()